package index

import (
	"bytes"
	"gitee.com/Cauchy_AQ/gokv/data"
	"sort"
	"sync"

	"github.com/google/btree"
)

// BTree 索引，封装 google 的 btree
type BTree struct {
	tree *btree.BTree
	lock *sync.RWMutex
}

// NewBTree 初始化 BTree 索引结构
func NewBTree() *BTree {
	return &BTree{
		tree: btree.New(32),
		lock: new(sync.RWMutex),
	}
}

// Put 向 BTree 中插入一条数据
func (b *BTree) Put(key []byte, pos *data.LogRecordPos) *data.LogRecordPos {
	item := &Item{key: key, pos: pos}
	b.lock.Lock()
	defer b.lock.Unlock()
	oldItem := b.tree.ReplaceOrInsert(item)
	if oldItem == nil {
		return nil
	}
	return oldItem.(*Item).pos
}

// Get 从 BTree 中获取索引为 key 的数据
func (b *BTree) Get(key []byte) *data.LogRecordPos {
	item := &Item{key: key}
	btreeItem := b.tree.Get(item)
	if btreeItem == nil {
		return nil
	}
	return btreeItem.(*Item).pos
}

// Delete 删除索引为 key 的数据
func (b *BTree) Delete(key []byte) (*data.LogRecordPos, bool) {
	item := &Item{key: key}
	b.lock.Lock()
	defer b.lock.Unlock()
	oldItem := b.tree.Delete(item)
	if oldItem == nil {
		return nil, false
	}
	return oldItem.(*Item).pos, true
}

// Iterator 获取迭代器
func (b *BTree) Iterator(reverse bool) Iterator {
	if b.tree == nil {
		return nil
	}
	b.lock.RLock()
	defer b.lock.RUnlock()

	return newBTreeIterator(b.tree, reverse)
}

// Size 获取 btree 结构数据大小
func (b *BTree) Size() int {
	return b.tree.Len()
}

// Close 关闭索引
func (b *BTree) Close() error {
	return nil
}

// btreeIterator BTree 索引迭代器
// 抽象这个迭代器是为了更好的支持需要的数据操作
type btreeIterator struct {
	currentIndex int     // 当前遍历的下标
	reverse      bool    // 是否反向遍历
	values       []*Item // 存储遍历出来的 key 和索引位置信息 data.LogRecordPos
}

// newBtreeIterator 新建 BTree 索引迭代器
func newBTreeIterator(tree *btree.BTree, reverse bool) *btreeIterator {
	var idx int
	values := make([]*Item, tree.Len())

	// 遍历数据回调，将btree遍历后的数据存储在自定义的迭代器结构values中
	saveValues := func(it btree.Item) bool {
		values[idx] = it.(*Item)
		idx++
		return true
	}

	if !reverse {
		// 正向遍历，并对数据调用 saveValues
		tree.Ascend(saveValues)
	} else {
		// 反向遍历，并对数据调用 saveValues
		tree.Descend(saveValues)
	}
	return &btreeIterator{
		currentIndex: 0,
		reverse:      reverse,
		values:       values,
	}
}

// Rewind 重新返回迭代起点
func (b *btreeIterator) Rewind() {
	b.currentIndex = 0
}

// Seek 根据传入 key 找到第一个大于（或小于）等于的目标 key，根据这个 key 开始遍历
func (b *btreeIterator) Seek(key []byte) {
	if b.reverse {
		// 反序，找到第一个小于等于 key 的位置
		b.currentIndex = sort.Search(len(b.values), func(i int) bool {
			return bytes.Compare(b.values[i].key, key) <= 0
		})
	} else {
		// 正序，找到第一个大于等于 key 的位置
		b.currentIndex = sort.Search(len(b.values), func(i int) bool {
			return bytes.Compare(b.values[i].key, key) >= 0
		})
	}
}

// Next 跳转到下一个 key
func (b *btreeIterator) Next() {
	b.currentIndex += 1
}

// Valid 是否有效，判断是否遍历完了 key
func (b *btreeIterator) Valid() bool {
	return b.currentIndex < len(b.values)
}

// Key 获取迭代器当前位置的 key
func (b *btreeIterator) Key() []byte {
	return b.values[b.currentIndex].key
}

// Value 获取当前迭代位置的索引位置信息
func (b *btreeIterator) Value() *data.LogRecordPos {
	return b.values[b.currentIndex].pos
}

// Close 关闭迭代器
func (b *btreeIterator) Close() {
	b.values = nil
}

/*
1. 接口去接指针或非指针的实例都可以，但接口不能采用指针
*/
